#include "cppv2.h"

#include "MxBase/Log/Log.h"
#include "MxBase/postprocess/include/ObjectPostProcessors/Yolov3PostProcess.h"
#include <algorithm>
#include <map>
#include "MxBase/Maths/FastMath.h"

// This option should be opened when using 200DK
// #define USE_200DK

using namespace MxBase;
using namespace std;
namespace{
    const uint32_t YUV_BYTE_NU = 3;
    const uint32_t YUV_BYTE_DE = 2;
    const uint32_t VPC_H_ALIGN = 2;
    const uint32_t YOLOV3_RESIZE = 416;
}

YoloV3Cppv2::YoloV3Cppv2(const V2Param &v2Param) 
{
    deviceId = v2Param.deviceId;
    std::string modelPath = v2Param.modelPath;
    APP_ERROR ret;

    // global init
    ret = MxInit();
    if (ret != APP_ERR_OK) 
    {
        LogError << "MxInit failed, ret=" << ret << ".";
    }

    // imageProcess init
    imageProcessorDptr = std::make_shared<MxBase::ImageProcessor>(deviceId); 
    if (imageProcessorDptr == nullptr)
    {
        LogError << "imageProcessorDptr nullptr";
    }
    

    // model init
    yoloV3Dptr = std::make_shared<MxBase::Model>(modelPath, deviceId);
    if (yoloV3Dptr == nullptr)
    {
        LogError << "yoloV3Dptr nullptr";
    }
};

APP_ERROR YoloV3Cppv2::ReadImage(const string &imgPath, Image &decodedImage)
{
    APP_ERROR ret;

#ifdef USE_200DK
    // 200DK only support this method!
    std::shared_ptr<uint8_t> dataPtr;
    uint32_t dataSize;
    Image imageData;
    ret = GetImage(imgPath, dataPtr, dataSize);
    if (ret != APP_ERR_OK) 
    {
        LogError << "Getimage failed, ret=" << ret;
        return ret;
    }
    ret = imageProcessorDptr->Decode(dataPtr, dataSize, decodedImage, ImageFormat::YUV_SP_420);
#else
    ret = imageProcessorDptr->Decode(imgPath, decodedImage, ImageFormat::YUV_SP_420);
#endif
    if (ret != APP_ERR_OK) 
    {
        LogError << "Decode failed, ret=" << ret;
        return ret;
    }
    LogInfo << "decodedImage.width is" << decodedImage.GetOriginalSize().width;
    LogInfo << "decodedImage.height is" << decodedImage.GetOriginalSize().height;
    LogInfo << "decodedImage.format is" << (int)(decodedImage.GetFormat());

    return APP_ERR_OK;
};

APP_ERROR YoloV3Cppv2::Resize(const Image& decodedImage, Image &resizeImage)
{
    APP_ERROR ret;

    // set size param
    Size resizeConfig(YOLOV3_RESIZE, YOLOV3_RESIZE);

    ret = imageProcessorDptr->Resize(decodedImage, resizeConfig, resizeImage, Interpolation::HUAWEI_HIGH_ORDER_FILTER);
    if (ret != APP_ERR_OK) 
    {
        LogError << "Resize failed, ret=" << ret;
        return ret;
    }

    // save resize image
    std::string path = "./resized_yolov3_416.jpg";
    ret = imageProcessorDptr->Encode(resizeImage, path);
    if (ret != APP_ERR_OK) 
    {
        LogError << "Encode failed, ret=" << ret;
        return ret;
    }
    return APP_ERR_OK;
};

APP_ERROR YoloV3Cppv2::Infer(Image &resizeImage, vector<Tensor>& yoloV3Outputs)
{
    APP_ERROR ret;
    // !move image to device!
    Tensor tensorImg = resizeImage.ConvertToTensor();
    ret = tensorImg.ToDevice(deviceId);
    if (ret != APP_ERR_OK) 
    {
        LogError << "ToDevice failed, ret=" << ret;
        return ret;
    }

    // make infer input
    std::vector<Tensor> yoloV3Inputs = {tensorImg};
    // do infer
    yoloV3Outputs = yoloV3Dptr->Infer(yoloV3Inputs);
    std::cout << "yoloV3Outputs len=" << yoloV3Outputs.size() << std::endl;

    // !move result to host!
    for (auto output : yoloV3Outputs)
    {
        output.ToHost();
    }
    
    return APP_ERR_OK;
};

APP_ERROR YoloV3Cppv2::tensorInfer(std::vector<MxBase::Tensor>& yoloV3Iutputs, std::vector<MxBase::Tensor>& yoloV3Outputs)
{
    // do infer
    yoloV3Outputs = yoloV3Dptr->Infer(yoloV3Iutputs);
    std::cout << "yoloV3Outputs len=" << yoloV3Outputs.size() << std::endl;

    // !move result to host!
    for (auto output : yoloV3Outputs)
    {
        output.ToHost();
    }
    return APP_ERR_OK;
};

void YoloV3Cppv2::YoloV3PostProcess(string& yoloV3ConfigPath, string& yoloV3LablePath,
                        vector<Tensor>& yoloV3Outputs, vector<Rect>& cropConfigVec)
{
    /// This should made by user! This func only show used with sdk`s yolov3 so lib.

    std::cout << "******YoloV3PostProcess******" << std::endl;
    // make yoloV3config map
    std::map<std::string, std::string> postConfig;

    postConfig.insert(pair<std::string, std::string>("postProcessConfigPath", yoloV3ConfigPath));
    postConfig.insert(pair<std::string, std::string>("lablePath", yoloV3LablePath));

    // init postProcess
    Yolov3PostProcess yolov3PostProcess;
    yolov3PostProcess.Init(postConfig);

    // make postProcess inputs
    vector<TensorBase> tensors;
    for (size_t i = 0; i < yoloV3Outputs.size(); i++)
    {
        MemoryData memoryData(yoloV3Outputs[i].GetData(), yoloV3Outputs[i].GetByteSize());
        TensorBase tensorBase(memoryData, true, yoloV3Outputs[i].GetShape(), TENSOR_DTYPE_INT32);
        tensors.push_back(tensorBase);
    }
    vector<vector<ObjectInfo>> objectInfos;

    auto shape = yoloV3Outputs[0].GetShape();
    MxBase::ResizedImageInfo imgInfo;
    imgInfo.widthOriginal = shape[1];
    imgInfo.heightOriginal = shape[0] * YUV_BYTE_DE / YUV_BYTE_NU;
    imgInfo.widthResize = YOLOV3_RESIZE;
    imgInfo.heightResize = YOLOV3_RESIZE;
    imgInfo.resizeType = MxBase::RESIZER_STRETCHING;
    std::vector<MxBase::ResizedImageInfo> imageInfoVec = {};
    imageInfoVec.push_back(imgInfo);
    
    // do postProcess
    yolov3PostProcess.Process(tensors, objectInfos, imageInfoVec);

    // print result
    std::cout << "Size of objectInfos is " << objectInfos.size() << std::endl;
    for (size_t i = 0; i < objectInfos.size(); i++)
    {
        std::cout << "objectInfo-" << i << " ,Size:"<< objectInfos[i].size() << std::endl;
        for (size_t j = 0; j < objectInfos[i].size(); j++)
        {
            std::cout << std::endl << "*****objectInfo-" << i << ":" << j << std::endl;
            std::cout << "x0 is " << objectInfos[i][j].x0 << std::endl;
            std::cout << "y0 is " << objectInfos[i][j].y0 << std::endl;
            std::cout << "x1 is " << objectInfos[i][j].x1 << std::endl;
            std::cout << "y1 is " << objectInfos[i][j].y1 << std::endl;
            std::cout << "confidence is " << objectInfos[i][j].confidence << std::endl;
            std::cout << "classId is " << objectInfos[i][j].classId << std::endl;
            std::cout << "className is " << objectInfos[i][j].className << std::endl;           
        }
    }
    
    // do image crop
    cropConfigVec.resize(objectInfos[0].size());
    for (size_t i = 0; i < objectInfos[0].size(); i++)
    {
        cropConfigVec[i].x0 = objectInfos[0][i].x0;
        cropConfigVec[i].y0 = objectInfos[0][i].y0;
        cropConfigVec[i].x1 = objectInfos[0][i].x1;
        cropConfigVec[i].y1 = objectInfos[0][i].y1;
    }
    
    std::cout << "******YoloV3PostProcess end******" << std::endl;
};

void YoloV3Cppv2::generalPostProcess(std::vector<MxBase::Tensor>& inferOutputs)
{
    // Modify this func in .cpp and .h to finish your postprocess, need add return
    size_t count = 0; // define result count or got by .Size()
    if (inferOutputs.empty())
    {
        std::cout << "result Infer failed with empty output..." << std::endl;
        return;
    }
    
    // convert tensor to cpp data. This sample convert first result tensor to float32 array
    float *castData = static_cast<float *>(inferOutputs[0].GetData());
    
    std::vector<float> result_0;
    for (size_t i = 0; i < count; i++)
    {
        result_0.push_back(castData[i]);
    }

    /// do compute

    return;
}

APP_ERROR YoloV3Cppv2::GetImage(const std::string &imgPath, std::shared_ptr<uint8_t>& dataPtr, uint32_t& dataSize)
{
    // Get image data to memory, this method can be substituted or designed by yourself!
    std::ifstream file;
    file.open(imgPath.c_str(), std::ios::binary);
    if (!file)
    {
        std::cout << "Invalid file." << std::endl;
        return APP_ERR_COMM_INVALID_PARAM;
    }
    std::stringstream buffer;
    buffer << file.rdbuf();
    std::string content = buffer.str();

    char* p = (char*) malloc(content.size());
    memcpy(p, content.data(), content.size());
    auto deleter = [](void *p) -> void {
        free(p);
        p = nullptr;
    };

    dataPtr.reset(static_cast<uint8_t*>((void*)(p)), deleter);
    dataSize = content.size();

    file.close();
    return APP_ERR_OK;
}